"""
Configuration classes for DIL_Frame models.

This module provides strongly typed configuration classes for all model components,
enhancing type safety, documentation, and maintainability. All configurations use
Python dataclasses for automatic initialization and serialization support.
"""

from dataclasses import dataclass, field
from typing import List, Optional, Union, Dict, Any
import torch


@dataclass
class BackboneConfig:
    """Base configuration for backbone models."""

    embed_dim: int = 768
    output_dim: int = 768
    dropout: float = 0.0


@dataclass
class ViTConfig(BackboneConfig):
    """Configuration for Vision Transformer models."""

    patch_size: int = 16
    img_size: int = 224
    in_chans: int = 3
    num_heads: int = 12
    depth: int = 12
    mlp_ratio: float = 4.0
    qkv_bias: bool = True
    attn_drop_rate: float = 0.0
    drop_path_rate: float = 0.0
    global_pool: bool = False
    use_cls_token: bool = True


@dataclass
class PEFTConfig:
    """Configuration for Parameter-Efficient Fine-Tuning methods."""

    method: str = "lora"  # lora, adapter
    bottleneck_dim: int = 64
    alpha: float = 1.0
    dropout: float = 0.1
    adapter_scalar: Union[float, str] = 1.0
    adapter_layernorm_option: str = "in"  # in, out, none


@dataclass
class ModelConfig:
    """Master configuration that combines all component configs."""

    backbone_type: str = "vit_base_patch16_224"
    backbone: ViTConfig = field(default_factory=ViTConfig)
    peft: Optional[PEFTConfig] = field(default_factory=PEFTConfig)
    init_cls: int = 10
    inc_cls: int = 10
    fc_func: str = "cosine"
    cam_visual: bool = False
    device_ids: List[int] = field(default_factory=lambda: [0])

    @property
    def device(self) -> torch.device:
        """Get the primary device for this model."""
        device_id = self.device_ids[0] if self.device_ids else 0
        if torch.cuda.is_available():
            return torch.device(f"cuda:{device_id}")
        return torch.device("cpu")

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> "ModelConfig":
        """Create a configuration from a dictionary."""
        # Convert dictionary representation to actual config instance
        backbone_cfg = ViTConfig(**config_dict.get("backbone", {}))
        peft_cfg = (
            PEFTConfig(**config_dict.get("peft", {})) if "peft" in config_dict else None
        )

        # Handle scalar values at the top level
        basic_fields = {
            k: v
            for k, v in config_dict.items()
            if k not in ["backbone", "peft", "device_ids"] and not k.startswith("_")
        }

        # Special handling for device
        device_ids = config_dict.get("device", [0])
        if isinstance(device_ids, (list, tuple)):
            device_ids = list(device_ids)
        else:
            device_ids = [device_ids]

        return cls(
            backbone=backbone_cfg, peft=peft_cfg, device_ids=device_ids, **basic_fields
        )
